#include "InvertedPendulum.h"
#include <iostream>
#include <QObject>
#include <QTimer>

InvertedPendulum::InvertedPendulum(double mass_cart, double mass_pole, double length_pole, double gravity)
    : m_cart_(mass_cart), m_pole_(mass_pole), l_pole_(length_pole), g_(gravity),
      x_(0.0), theta_(0.1), x_dot_(0.0), theta_dot_(0.0) {
    // 初始化状态空间矩阵 A
    A = Eigen::MatrixXd(4, 4);
    A << 0, 1, 0, 0,
         0, 0, -m_pole_ * g_ / m_cart_, 0,
         0, 0, 0, 1,
         0, 0, (m_cart_ + m_pole_) * g_ / (m_cart_ * l_pole_), 0;

    // 初始化控制输入矩阵 B
    B = Eigen::MatrixXd(4, 1);
    B << 0,
         1 / m_cart_,
         0,
        -1 / (m_cart_ * l_pole_);
}

double InvertedPendulum::GetPosition() const {
    return x_;
}

double InvertedPendulum::GetAngle() const {
    return theta_;
}

Eigen::VectorXd InvertedPendulum::LQRControl() {
    // LQR权重矩阵 Q 和 R
    Eigen::MatrixXd Q(4, 4);
    Q << 1, 0, 0, 0,
         0, 1, 0, 0,
         0, 0, 10, 0,
         0, 0, 0, 10;

    Eigen::MatrixXd R(1, 1);
    R << 1;

    // 使用LQR计算控制增益矩阵 K
    Eigen::MatrixXd K = Eigen::MatrixXd::Zero(1, 4);
    Eigen::MatrixXd P = Q;

    for (int i = 0; i < 100; ++i) {
        K = (R + B.transpose() * P * B).inverse() * B.transpose() * P * A;
        P = Q + A.transpose() * P * (A - B * K);
    }

    // 计算控制输入 u = -K * x
    Eigen::VectorXd x(4);
    x << x_, x_dot_, theta_, theta_dot_;
    Eigen::VectorXd control_input = -K * x;

    return control_input;
}

void InvertedPendulum::ControlLoop() {
    Eigen::VectorXd control_input = LQRControl();

    // 在这里根据控制输入更新倒立摆的状态
    // 更新倒立摆状态，这里简单示范
    Eigen::VectorXd x(4);
    x << x_, x_dot_, theta_, theta_dot_;
    Eigen::VectorXd x_new = A * x + B * control_input;
    x_ = x_new(0);
    x_dot_ = x_new(1);
    theta_ = x_new(2);
    theta_dot_ = x_new(3);
}
void InvertedPendulum::ApplyDisturbance(double dragStartPosx,double dragStartPosy,double forcex,double forcey){
    theta_ += forcex*forcey*100;
}
